from __future__ import annotations

import logging
import time
from datetime import datetime, date, timedelta
from typing import Iterable, Dict, Any, Optional, Tuple
import pandas as pd
from chinese_calendar import is_workday 

from quantcore.db.mongo import Mongo

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 常用指数代码列表
COMMON_INDICES = {
    "000001": "上证指数",
    "000016": "上证50", 
    "000300": "沪深300",
    "000905": "中证500",
    "000852": "中证1000",
    "399001": "深证成指",
    "399006": "创业板指",
    "000688": "科创50",
    "000010": "上证180",
    "399004": "深证100",
    "399005": "中小100",
    "899050": "北证50"
}


class MarketDataCollector:
    """AKShare行情数据采集器，实现FR-DM-001需求。"""

    def __init__(self, request_interval: float = 3.0):
        """
        初始化数据采集器
        
        Args:
            request_interval: 请求间隔时间（秒），建议≥3秒避免被接口拒绝
        """
        self.db = Mongo.db()
        self.request_interval = max(request_interval, 1.0)  # 最小间隔1秒
        self.last_request_time = 0.0

        
    def check_data_freshness(self, collection_name: str = "price_daily") -> Tuple[Optional[date], Optional[date]]:
        """检查数据最新性，返回需要更新的日期范围"""
        try:
            latest_record = self.db[collection_name].find_one({}, sort=[('trade_date', -1)])
            
            if latest_record and 'trade_date' in latest_record:
                db_latest = latest_record['trade_date'].date()
            else:
                db_latest = date(2024, 1, 1)
            
            # 获取最近的交易日期，而不是当前日期
            latest_trading_date = self._get_latest_trading_date()
            
            if latest_trading_date and db_latest < latest_trading_date:
                start_date = db_latest + timedelta(days=1)
                end_date = latest_trading_date
                logger.info(f"数据需要更新：{start_date} 到 {end_date}")
                return start_date, end_date
            else:
                logger.info("数据已是最新，无需更新")
                return None, None
                
        except Exception as e:
            logger.error(f"检查数据最新性失败: {e}")
            return None, None
    
    def _get_latest_trading_date(self) -> Optional[date]:
        """获取最近的交易日期"""
        try:
            current_date = datetime.now().date()
            logger.debug(f"当前日期: {current_date}")
            
            # 如果当前是交易日，返回当前日期
            if self._is_trading_day(current_date):
                logger.info(f"当前日期 {current_date} 是交易日")
                return current_date
            else:
                # 向前查找最近的交易日
                logger.debug("当前不是交易日，向前查找最近交易日")
                for i in range(1, 8):  # 最多向前查找7天
                    check_date = current_date - timedelta(days=i)
                    if self._is_trading_day(check_date):
                        logger.info(f"向前查找到最近交易日: {check_date}")
                        return check_date
                
                logger.warning("向前查找7天仍未找到交易日")
                return None
                    
        except Exception as e:
            logger.error(f"获取最近交易日期失败: {e}")
            return None
    

    
    def _control_request_rate(self):
        """控制请求频率，避免被接口拒绝"""
        current_time = time.time()
        time_since_last_request = current_time - self.last_request_time
        
        if time_since_last_request < self.request_interval:
            sleep_time = self.request_interval - time_since_last_request
            logger.debug(f"请求间隔控制：等待 {sleep_time:.2f} 秒")
            time.sleep(sleep_time)
        
        self.last_request_time = time.time()
    
    def _is_trading_day(self, check_date: date) -> bool:
        """判断是否为交易日（使用chinese_calendar）"""
        try:
            # 使用chinese_calendar判断是否为工作日（交易日）
            is_trading = is_workday(check_date)
            logger.debug(f"日期 {check_date} 是否为交易日: {is_trading}")
            return is_trading
            
        except ImportError:
            logger.warning("未安装chinese_calendar，使用周末判断作为备选")
            # 备选方案：简单周末判断
            if check_date.weekday() >= 5:
                return False
            return True
        except Exception as e:
            logger.warning(f"交易日判断失败，使用周末判断作为备选: {e}")
            # 备选方案：简单周末判断
            if check_date.weekday() >= 5:
                return False
            return True
    

    
    def fetch_stock_list(self) -> pd.DataFrame:
        """获取全量A股股票代码列表"""
        try:
            self._control_request_rate()  # 控制请求频率
            import akshare as ak
            stock_list = ak.stock_info_a_code_name()
            logger.info(f"成功获取股票列表，共{len(stock_list)}只股票")
            return stock_list
        except ImportError:
            logger.error("未安装akshare，请先安装：pip install akshare")
            return pd.DataFrame()
        except Exception as e:
            logger.error(f"获取股票列表失败: {e}")
            return pd.DataFrame()
    
    def fetch_daily_prices(self, symbol: str, start_date: date, end_date: date, 
                          adjust: str = "qfq") -> Optional[pd.DataFrame]:
        """获取个股日线数据"""
        try:
            self._control_request_rate()  # 控制请求频率
            import akshare as ak
            
            df = ak.stock_zh_a_hist(
                symbol=symbol,
                period="daily",
                start_date=start_date.strftime("%Y%m%d"),
                end_date=end_date.strftime("%Y%m%d"),
                adjust=adjust
            )
            
            if df.empty:
                logger.warning(f"股票{symbol}在{start_date}到{end_date}期间无数据")
                return None
                
            df = self._clean_daily_data(df, symbol)
            return df
            
        except Exception as e:
            logger.error(f"获取股票{symbol}日线数据失败: {e}")
            return None
    
    def fetch_index_daily(self, symbol: str, start_date: date, end_date: date) -> Optional[pd.DataFrame]:
        """获取指数日线数据"""
        try:
            self._control_request_rate()  # 控制请求频率
            import akshare as ak
            
            df = ak.index_zh_a_hist(
                symbol=symbol,
                period="daily",
                start_date=start_date.strftime("%Y%m%d"),
                end_date=end_date.strftime("%Y%m%d")
            )
            
            if df.empty:
                logger.warning(f"指数{symbol}在{start_date}到{end_date}期间无数据")
                return None
                
            df = self._clean_index_data(df, symbol)
            return df
            
        except Exception as e:
            logger.error(f"获取指数{symbol}日线数据失败: {e}")
            return None
    
    def _clean_daily_data(self, df: pd.DataFrame, symbol: str) -> pd.DataFrame:
        """清洗个股日线数据"""
        try:
            column_mapping = {
                '日期': 'trade_date',
                '开盘': 'open',
                '最高': 'high', 
                '最低': 'low',
                '收盘': 'close',
                '成交量': 'volume',
                '成交额': 'turnover',
                '振幅': 'amplitude',
                '涨跌幅': 'change_pct',
                '涨跌额': 'change_amount',
                '换手率': 'turnover_rate'
            }
            
            df = df.rename(columns=column_mapping)
            df['trade_date'] = pd.to_datetime(df['trade_date'])
            
            numeric_columns = ['open', 'high', 'low', 'close', 'volume', 'turnover', 'turnover_rate']
            for col in numeric_columns:
                if col in df.columns:
                    df[col] = pd.to_numeric(df[col], errors='coerce')
            
            df['symbol'] = symbol
            df['adj_factor'] = 1.0
            df['is_suspended'] = False
            df['limit_status'] = 0
            
            return df
            
        except Exception as e:
            logger.error(f"清洗股票{symbol}数据失败: {e}")
            return df
    
    def _clean_index_data(self, df: pd.DataFrame, symbol: str) -> pd.DataFrame:
        """清洗指数日线数据"""
        try:
            column_mapping = {
                '日期': 'trade_date',
                '开盘': 'open',
                '最高': 'high',
                '最低': 'low', 
                '收盘': 'close',
                '成交量': 'volume',
                '成交额': 'turnover'
            }
            
            df = df.rename(columns=column_mapping)
            df['trade_date'] = pd.to_datetime(df['trade_date'])
            
            numeric_columns = ['open', 'high', 'low', 'close', 'volume', 'turnover']
            for col in numeric_columns:
                if col in df.columns:
                    df[col] = pd.to_numeric(df[col], errors='coerce')
            
            df['symbol'] = symbol
            df['index_level'] = df['close']
            df['free_float'] = 0.0
            df['free_float_mv'] = 0.0
            
            return df
            
        except Exception as e:
            logger.error(f"清洗指数{symbol}数据失败: {e}")
            return df
    
    def save_to_mongodb(self, df: pd.DataFrame, collection_name: str, symbol: str) -> bool:
        """将数据保存到MongoDB"""
        try:
            if df.empty:
                logger.warning(f"数据为空，跳过保存到{collection_name}")
                return False
            
            records = df.to_dict('records')
            collection = self.db[collection_name]
            
            for record in records:
                if 'trade_date' in record and 'symbol' in record:
                    query = {
                        'trade_date': record['trade_date'],
                        'symbol': record['symbol']
                    }
                    
                    collection.update_one(
                        query,
                        {'$set': record},
                        upsert=True
                    )
            
            logger.info(f"成功保存{len(records)}条{symbol}数据到{collection_name}")
            return True
            
        except Exception as e:
            logger.error(f"保存{symbol}数据到{collection_name}失败: {e}")
            return False
    
    def fetch_market_data(self) -> bool:
        """带数据最新性检查的市场数据采集主函数"""
        try:
            start_date, end_date = self.check_data_freshness()
            
            if not start_date or not end_date:
                logger.info("数据已是最新，无需更新")
                return True
            
            logger.info(f"开始更新数据：{start_date} 到 {end_date}")
            logger.info(f"请求间隔设置：{self.request_interval}秒（建议≥3秒避免被接口拒绝）")
            
            stock_list = self.fetch_stock_list()
            if stock_list.empty:
                logger.error("获取股票列表失败")
                return False
            
            success_count = 0
            total_count = len(stock_list)
            
            for idx, row in stock_list.iterrows():
                try:
                    symbol = row.get('code', row.get('symbol', ''))
                    if not symbol:
                        continue
                        
                    df = self.fetch_daily_prices(symbol, start_date, end_date)
                    if df is not None and not df.empty:
                        if self.save_to_mongodb(df, "price_daily", symbol):
                            success_count += 1
                    
                    if (idx + 1) % 100 == 0:
                        logger.info(f"股票数据采集进度: {idx + 1}/{total_count}，请求间隔: {self.request_interval}秒")
                        
                except Exception as e:
                    logger.error(f"处理股票{row.get('code', 'unknown')}时出错: {e}")
                    continue
            
            index_success_count = 0
            logger.info(f"开始采集指数数据，共{len(COMMON_INDICES)}个指数")
            for symbol, name in COMMON_INDICES.items():
                try:
                    df = self.fetch_index_daily(symbol, start_date, end_date)
                    if df is not None and not df.empty:
                        if self.save_to_mongodb(df, "index_daily", symbol):
                            index_success_count += 1
                            logger.debug(f"成功采集指数{symbol}({name})数据")
                            
                except Exception as e:
                    logger.error(f"处理指数{symbol}({name})时出错: {e}")
                    continue
            
            logger.info(f"数据更新完成：股票{success_count}/{total_count}，指数{index_success_count}/{len(COMMON_INDICES)}")
            return success_count > 0
            
        except Exception as e:
            logger.error(f"市场数据采集失败: {e}")
            return False


class FundamentalsCollector:
    """基本面采集占位"""

    def fetch_quarterly_financials(self) -> Iterable[Dict[str, Any]]:
        raise NotImplementedError

